import os
import pandas as pd
import numpy as np


# # 读取表格
# dfs = []
# input_dir = "./13.tree_annotation_to_table"
# for index, file in enumerate(os.listdir(input_dir)):
#     print(index, "\t", file)
#     file_path = os.path.join(input_dir, file)
#     df = pd.read_excel(file_path, keep_default_na=False)
#     dfs.append(df)
# df = pd.concat(dfs)

from multiprocessing import Pool


dir = "./13.tree_annotation_to_table"
files = os.listdir(dir)

def read_df(file_path, name):
    print(file_path)
    df = pd.read_excel(file_path, keep_default_na=False)
    return df

tasks = []
for f in files:
    tasks.append((os.path.join(dir,f),""))

with Pool(250) as pool:
    dfs = pool.starmap(read_df, tasks)

df = pd.concat(dfs)

# 替换空字符串为 NaN
df['gene_local_density'] = df['gene_local_density'].replace('', np.nan)
df['gene_local_density'] = df['gene_local_density'].fillna(method='ffill')
title = df.columns.tolist()
print(title)
# print(df)

# 读取物种相关名字或id表
#species_info_df = pd.read_csv("./cell_species.csv", header=None)
species_info_df = pd.read_csv("./20250814_cell_species.csv", header=None)   #物种信息表格使用最新的表格
species_dict = {}
for row in species_info_df.values.tolist():
    species_dict[row[-1]] = row[0:2]

# for k,v in species_dict.items():
#     print(k,v)

[
    'species', 'genome_gene_number', 'genome_length', 'genome_gene_density',
    'chromosome/scaffold/contif', 'cluster_name', 'transcript_id', 'gene_id',
    'start', 'end', 'strand', 'gene_local_density', 'cluster_gene_density',
    'min', 'quartile_lower', 'quartile_median', 'quartile_upper', 'max',
    'gene_description', 'gog_category', 'pfam', 'main_category',
    'summary_feature', 'level_1', 'level_2', 'level_3', 'level_4', 'level_5', 'level_6'
]

# -------------------生成 gene_info 信息表--------------------------------------
gene_info_data = []
for index, row in df.iterrows():
    try:
        genome = row["species"]
        species_id = species_dict[genome][0]
        chromosome = row["chromosome/scaffold/contif"]
        cluster_name = row["cluster_name"]
        transcript_id = row["transcript_id"]
        gene_id = row["gene_id"]
        start = row["start"]
        end = row["end"]
        strand = row["strand"]
        density = [row["gene_local_density"],row["cluster_gene_density"]]
        pfam = row["pfam"]
        gene_description = row["gene_description"]
        cog_category = row["gog_category"]
        main_category = row["main_category"]
        summary_feature = row["summary_feature"]
        level_1 = row["level_1"]
        level_2 = row["level_2"]
        level_3 = row["level_3"]
        level_4 = row["level_4"]
        level_5 = row["level_5"]
        level_6 = row["level_6"]


        text = (
            [genome, species_id] + [chromosome, cluster_name, transcript_id] +
            [gene_id, start, end, strand] +  density +
            [pfam, gene_description, cog_category, main_category, summary_feature] +
            [level_1, level_2, level_3, level_4, level_5, level_6]
        )
        gene_info_data.append([(" " if str(x) == "nan" else x) for x in text])
    except Exception as e:
        print(e)
        # print(row)


title = [
    "genome", "genome_id", "chromosome", "cluster_name", "transcript_id",
    "gene_id", "start", "end", "strand", "local_gene_density", "cluster_gene_density",
    "pfam", "gene_description", "cog_category", "main_category", "summary_feature",
    "level_1", "level_2", "level_3", "level_4", "level_5", "level_6"
]

output_dir = "17.gene_info"
os.makedirs(output_dir, exist_ok=True)
df_out_path = os.path.join(output_dir, "gene_info.csv")
df = pd.DataFrame(gene_info_data, columns=title)
df.to_csv(df_out_path, index=False)

sql_out_path = os.path.join(output_dir, "gene_info.sql")
sql = (
    "INSERT INTO gene_info " + str(tuple(title)).replace("'", "`") + " VALUES " +
    ",\n".join([str(tuple(x)) for x in gene_info_data]) + ";"
)
with open(sql_out_path, "w", encoding="utf-8") as f:
    f.write(sql)

